#!/usr/bin/env bash
#
# PSI performance comparison benchmark
# Requires: root privileges (for dropping caches), stress-ng, fio, vmstat, iostat
#
# This script runs four workloads (cpu, memory, io, mixed) twice (implicitly),
# depending on whether the kernel was booted with psi=0 or not.
# Results are stored under ./psi_test_results.
#
# IO test has been rewritten to use fio with deterministic random read/write
# pattern using a fixed random seed to improve reproducibility.

set -euo pipefail

########################
# Configuration
########################
DURATION=${DURATION:-300}                 # Duration (seconds) per workload
OUTPUT_DIR=${OUTPUT_DIR:-"./psi_test_results"}
IO_WORK_DIR=${IO_WORK_DIR:-"/test"}       # Directory on a filesystem with enough free space
IO_FILE=${IO_FILE:-"${IO_WORK_DIR}/psi_randrw_testfile.dat"}
IO_FILE_SIZE=${IO_FILE_SIZE:-"8G"}        # Size of the fio test file
IO_BS=${IO_BS:-"128k"}                    # Block size
IO_RWMIX_READ=${IO_RWMIX_READ:-50}        # Read percentage in mixed randrw
IO_IODEPTH=${IO_IODEPTH:-64}
IO_NUMJOBS=${IO_NUMJOBS:-$(nproc)}        # Concurrency (jobs)
IO_RANDOM_SEED=${IO_RANDOM_SEED:-123456}  # Fixed random seed
IO_ENGINE=${IO_ENGINE:-"libaio"}          # Can switch to "psync" if libaio not present
FIO_EXTRA=${FIO_EXTRA:-""}                # Extra overrides if needed

# Drop caches before IO/mixed runs (1=yes / 0=no)
#DROP_CACHES_ON_IO=${DROP_CACHES_ON_IO:-1}
DROP_CACHES_ON_IO=1

mkdir -p "${OUTPUT_DIR}"

########################
# Functions
########################
check_deps() {
    local missing=0
    for cmd in stress-ng fio vmstat iostat; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            echo "ERROR: Required command '$cmd' not found in PATH."
            missing=1
        fi
    done
    if [[ $missing -eq 1 ]]; then
        echo "Please install missing dependencies and re-run."
        exit 1
    fi
}

require_root() {
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: This script must be run as root (needed for cache dropping)."
        exit 1
    fi
}

get_psi_state() {
    if grep -qw "psi=0" /proc/cmdline; then
        echo "disabled"
    else
        echo "enabled"
    fi
}

prepare_io_file() {
    if [[ ! -d "$IO_WORK_DIR" ]]; then
        echo "ERROR: IO working directory '$IO_WORK_DIR' does not exist."
        exit 1
    fi
    if [[ ! -f "$IO_FILE" ]]; then
        echo "Preparing fio data file: $IO_FILE ($IO_FILE_SIZE)"
        fio --name=prepare_io_file \
            --filename="${IO_FILE}" \
            --size="${IO_FILE_SIZE}" \
            --rw=write \
            --bs=1M \
            --ioengine="${IO_ENGINE}" \
            --iodepth=32 \
            --numjobs=1 \
            --direct=1 \
            --group_reporting=1 \
            --fallocate=truncate \
            --loops=1 \
            --randrepeat=1 \
            --randseed=${IO_RANDOM_SEED} \
            --time_based=0 >/dev/null 2>&1
        sync
    else
        echo "Reusing existing fio data file: $IO_FILE"
    fi
}

drop_caches_if_requested() {
    if [[ "${DROP_CACHES_ON_IO}" -eq 1 ]]; then
        echo "Dropping filesystem caches..."
        sync
        echo 3 > /proc/sys/vm/drop_caches
    fi
}

run_cpu_test() {
    local output="$1"
    stress-ng --cpu $(( $(nproc) * 4 )) \
              -t "${DURATION}" \
              --metrics-brief 2>&1 | tee -a "${output}"
}

run_memory_test() {
    local output="$1"
    # Adjust --vm-bytes if 24G is too large for your system
    stress-ng --vm $(( $(nproc) * 2 )) \
              --vm-bytes 24G \
              -t "${DURATION}" \
              --metrics-brief 2>&1 | tee -a "${output}"
}

run_io_test() {
    local output="$1"
    prepare_io_file
    drop_caches_if_requested

    # Use time_based test with deterministic random pattern
    # We emit both text (tee) and JSON (for analysis).
    local json_out="${output}.fio.json"

    fio --name=psi_randrw \
        --filename="${IO_FILE}" \
        --size="${IO_FILE_SIZE}" \
        --rw=randrw \
        --rwmixread="${IO_RWMIX_READ}" \
        --bs="${IO_BS}" \
        --ioengine="${IO_ENGINE}" \
        --iodepth="${IO_IODEPTH}" \
        --numjobs="${IO_NUMJOBS}" \
        --time_based=1 \
        --runtime="${DURATION}" \
        --group_reporting=1 \
        --direct=1 \
        --randrepeat=1 \
        --randseed="${IO_RANDOM_SEED}" \
        --random_generator=tausworthe64 \
        --invalidate=1 \
        --exitall_on_error=1 \
        --output-format=json \
        --output="${json_out}" \
        ${FIO_EXTRA} 2>&1 | tee -a "${output}"

    echo "fio JSON results saved to: ${json_out}" | tee -a "${output}"
}

run_mixed_test() {
    local output="$1"
    drop_caches_if_requested
    # Mixed stress: CPU, VM, and IO (stress-ng style).
    # You may later convert this to a combined fio + stress-ng scenario if desired.
    stress-ng --cpu "$(nproc)" \
              --iomix "$(nproc)" --iomix-bytes 1G --temp-path "${IO_WORK_DIR}" \
              --vm "$(nproc)" --vm-bytes 12G \
              -t "${DURATION}" \
              --metrics-brief 2>&1 | tee -a "${output}"
}

start_monitors() {
    local base="$1"
    vmstat 1 "${DURATION}" > "${base}.vmstat" &
    local vm_pid=$!
    iostat -dxm 1 "${DURATION}" > "${base}.iostat" &
    local io_pid=$!
    echo "${vm_pid} ${io_pid}"
}

run_test() {
    local test_name=$1
    local psi_state
    psi_state=$(get_psi_state)
    local output="${OUTPUT_DIR}/${test_name}_${psi_state}.txt"

    echo "====== Starting test: ${test_name} (PSI ${psi_state}) ======"
    echo "Parameters: duration=${DURATION}s | timestamp=$(date -Iseconds)" | tee -a "${output}"

    # Monitors
    pids=$(start_monitors "${output}")
    set -- ${pids}
    local vm_pid=$1
    local io_pid=$2

    case "${test_name}" in
        cpu)    run_cpu_test "${output}" ;;
        memory) run_memory_test "${output}" ;;
        io)     run_io_test "${output}" ;;
        mixed)  run_mixed_test "${output}" ;;
        *) echo "Unknown test type: ${test_name}" ; kill "${vm_pid}" "${io_pid}" 2>/dev/null || true ; return 1 ;;
    esac

    # Ensure monitors finish
    wait "${vm_pid}" "${io_pid}" || true

    echo -e "====== Test complete: ${test_name} ======\n" | tee -a "${output}"
}

main() {
    require_root
    check_deps

    echo "Detected PSI state: $(get_psi_state)"
    echo "Results directory: ${OUTPUT_DIR}"
    echo "Test duration per workload: ${DURATION}s"
    echo "IO test file: ${IO_FILE} (size=${IO_FILE_SIZE})"
    echo

    for test_type in cpu memory io mixed; do
        run_test "${test_type}"
    done

    echo "All tests finished. Results stored in: ${OUTPUT_DIR}"
}

main "$@"
